package com.dawn.framework.web.core.filter;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * @description:
 * @author: cnsu-cmh
 * @time: 2024/11/7 16:58
 */
public class ChangerIpResponseWrapper extends HttpServletResponseWrapper {

    private ServletOutputStream filterOutput;

    private ByteArrayOutputStream output;

    public ChangerIpResponseWrapper(HttpServletResponse response) {
        super(response);
        output = new ByteArrayOutputStream();
        filterOutput = new WrapperOutputStream(output);
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return filterOutput;
    }

    @Override
    public void flushBuffer() throws IOException {
        if (filterOutput != null) {
            filterOutput.flush();
        }
    }

    public byte[] getContent() throws IOException {
        flushBuffer();
        return output.toByteArray();
    }

    static class WrapperOutputStream extends ServletOutputStream {

        private final ByteArrayOutputStream bos;

        public WrapperOutputStream(ByteArrayOutputStream bos) {
            this.bos = bos;
        }

        @Override
        public void write(int b) {
            bos.write(b);
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener arg0) {
        }
    }

}

